"""
一个简单的demo，实现了角色扮演的基本功能

依赖：
pyjwt
requests
streamlit
zhipuai
python-dotenv

运行方式：
```bash
streamlit run role_play.py
```
"""
import time
import os
import json
from datetime import datetime
import random
import itertools
from typing import Iterator, Optional

import streamlit as st
from dotenv import load_dotenv
# 通过.env文件设置环境变量
# reference: https://github.com/theskumar/python-dotenv
load_dotenv()

import api
from api import generate_chat_scene_prompt, generate_role_appearance, get_characterglm_response, generate_cogview_image
from data_types import TextMsg, ImageMsg, TextMsgList, MsgList, filter_text_msg

st.set_page_config(page_title="CharacterGLM API Demo", page_icon="🤖", layout="wide")

def update_api_key(key: Optional[str] = None):
    if key:
        api.API_KEY = key

# 设置API KEY
api_key = st.sidebar.text_input("API_KEY", value=os.getenv("API_KEY", ""), key="API_KEY", type="password", on_change=update_api_key)
update_api_key(api_key)


# 初始化
if "history" not in st.session_state:
    st.session_state["history"] = []

def init_session():
    st.session_state["history"] = []


# 4个输入框，设置meta的4个字段
meta_labels = {
    "bot_name": "角色名",
    "bot_info": "角色人设",
    "user_name": "用户名", 
    "user_info": "用户人设"
}

liXiaoYao_info = "李逍遥，19岁，蜀山仙剑派，李逍遥是李三思的儿子，李寒空的后人；出生于盛渔村，在李家客栈当店小二。后经历种种，拜入蜀山仙剑派，还担任了第二十七代掌门；而后退位，成为蜀山七圣之一。悟性极高，聪明绝顶，富有强烈的正义感。少年时略滑头带点鬼灵精 。即使是小渔村中的店小二，但在客栈中英雄好汉见得多了。李逍遥的梦想就是成为“大侠”，在他心目中，侠就是像父亲李三思那样仗剑江湖、惩恶扬善，闯出一番声名，令世人尊敬的人。命带桃花，初闯江湖即与三名女主角产生一段难分难解的四角恋情。侠者，仁义之人也。李逍遥一路所行从未违背“侠”的准则，自然当得“侠”字美称。苗疆一场大劫后回到蜀山，静心练剑，年纪轻轻便完成了当时仙剑派掌门人“剑圣”独孤宇云所交付的七大难事与四十九项考验，当上了仙剑派第二十七代掌门。从此御剑行遍天下，行侠仗义、斩妖除魔，而且武功之高已入化境，武林人士都尊崇为天下武功第一的当代剑侠，但是心中念念不忘的就是寻找灵儿的下落，以及照顾月如直至完全康复的执着。"
aNu_info = "阿奴，性别女，22岁，白苗族少主、巫月神殿教主、圣姑，白苗族的少主。外表一派天真烂漫，娇俏可爱，施用巫术毒蛊的手段却叫人毛骨悚然。对李逍遥一往情深，但逍遥已有所爱，只得默默伴随一旁。后成为巫月神教的教主，担任新一代圣姑，人称“海棠夫人”。收李逍遥的外孙女小蛮为徒。少女时期个性刁钻伶俐，活泼可爱，年纪小却很精明；口快心直，却不失俏皮；虽然外表天真烂漫，但巫术毒蛊却运用自如。小小年纪，身法、巫术、武功已是不俗，更有平息黑白苗矛盾的心志。一路帮助李逍遥和赵灵儿与以拜月为首的邪恶势力做斗争。在感情里，阿奴无私无悔，她对李逍遥的这份情愫一直很安静，也很伟大，尽己所能帮助逍遥和灵儿，为所爱之人的默默付出，不求任何回报，甚至为爱斩断姻缘。阿奴成为巫月神殿掌教后，改名“海棠”，带领族人守护女娲神殿，精通苗族蛊医巫三术，在苗族中的地位倍受尊崇，从当年的少女成长为后来的成熟女性，并收养了一贫（李逍遥）的外孙女小蛮为徒，自少女时期便一直深爱着一贫，时隔多年两人却渐渐成了彼此最熟悉的陌生人。"

aNu_character_meta = {
        "user_name": "李逍遥",
        "user_info": liXiaoYao_info,
        "bot_info": aNu_info,
        "bot_name": "阿奴"
    }

liXiaoYao_character_meta = {
        "user_name": "阿奴",
        "user_info": aNu_info,
        "bot_info": liXiaoYao_info,
        "bot_name": "李逍遥"
    }

# 2x2 layout
with st.container():
    col1, col2 = st.columns(2)
    with col1:
        st.write("""
            # 阿奴
            阿奴，性别女，22岁，白苗族少主、巫月神殿教主、圣姑，白苗族的少主。外表一派天真烂漫，娇俏可爱，施用巫术毒蛊的手段却叫人毛骨悚然。对李逍遥一往情深，但逍遥已有所爱，只得默默伴随一旁。后成为巫月神教的教主，担任新一代圣姑，人称“海棠夫人”。收李逍遥的外孙女小蛮为徒。少女时期个性刁钻伶俐，活泼可爱，年纪小却很精明；口快心直，却不失俏皮；虽然外表天真烂漫，但巫术毒蛊却运用自如。小小年纪，身法、巫术、武功已是不俗，更有平息黑白苗矛盾的心志。一路帮助李逍遥和赵灵儿与以拜月为首的邪恶势力做斗争。在感情里，阿奴无私无悔，她对李逍遥的这份情愫一直很安静，也很伟大，尽己所能帮助逍遥和灵儿，为所爱之人的默默付出，不求任何回报，甚至为爱斩断姻缘。阿奴成为巫月神殿掌教后，改名“海棠”，带领族人守护女娲神殿，精通苗族蛊医巫三术，在苗族中的地位倍受尊崇，从当年的少女成长为后来的成熟女性，并收养了一贫（李逍遥）的外孙女小蛮为徒，自少女时期便一直深爱着一贫，时隔多年两人却渐渐成了彼此最熟悉的陌生人。
        """)
        
    with col2:
        st.write("""
            # 李逍遥
            李逍遥，19岁，蜀山仙剑派，李逍遥是李三思的儿子，李寒空的后人；出生于盛渔村，在李家客栈当店小二。后经历种种，拜入蜀山仙剑派，还担任了第二十七代掌门；而后退位，成为蜀山七圣之一。悟性极高，聪明绝顶，富有强烈的正义感。少年时略滑头带点鬼灵精 。即使是小渔村中的店小二，但在客栈中英雄好汉见得多了。李逍遥的梦想就是成为“大侠”，在他心目中，侠就是像父亲李三思那样仗剑江湖、惩恶扬善，闯出一番声名，令世人尊敬的人。命带桃花，初闯江湖即与三名女主角产生一段难分难解的四角恋情。侠者，仁义之人也。李逍遥一路所行从未违背“侠”的准则，自然当得“侠”字美称。苗疆一场大劫后回到蜀山，静心练剑，年纪轻轻便完成了当时仙剑派掌门人“剑圣”独孤宇云所交付的七大难事与四十九项考验，当上了仙剑派第二十七代掌门。从此御剑行遍天下，行侠仗义、斩妖除魔，而且武功之高已入化境，武林人士都尊崇为天下武功第一的当代剑侠，但是心中念念不忘的就是寻找灵儿的下落，以及照顾月如直至完全康复的执着。
        """)

button_labels = {
    "start_chat": "开始对话",
    "go_on": "继续对话",
    "clear_history": "清空对话历史",
    "save_chat": "保存对话"
}

# 展示对话历史
for msg in st.session_state["history"]:
    if msg["role"] == "user":
        with st.chat_message(name="user", avatar="user"):
            st.markdown(msg["content"])
    elif msg["role"] == "assistant":
        with st.chat_message(name="assistant", avatar="assistant"):
            st.markdown(msg["content"])

    else:
        raise Exception("Invalid role")


with st.chat_message(name="LiXIaoYao", avatar="user"):
    input_placeholder = st.empty()
with st.chat_message(name="ANu", avatar="assistant"):
    message_placeholder = st.empty()

def save_chat_history():
    # 获取聊天历史
    chat_history = st.session_state['history']
    # 获取当前日期和时间的字符串
    current_time = datetime.now().strftime('%Y%m%d_%H%M%S')
    # 生成一个随机数
    random_number = random.randint(1000, 9999)
    # 设置保存的文件名，包括日期和随机数
    filename = f'chat_history_{current_time}_{random_number}.json'
    # 打开文件，准备写入
    with open(filename, 'w', encoding='utf-8') as f:
        # 将聊天历史转换成JSON格式，并写入文件
        json.dump(chat_history, f, ensure_ascii=False, indent=2)
    # 通知用户文件已保存
    st.success('聊天历史已保存到: ' + filename)

def output_stream_response(response_stream: Iterator[str], placeholder):
    content = ""
    for content in itertools.accumulate(response_stream):
        placeholder.markdown(content)
    return content

def characterglm_ANu():
    if st.session_state["history"]:
        query = last_message = st.session_state["history"][-1].get("content")
    else:
        query = "（头晕晕的，疑惑的睁开眼，看着阿奴）我怎么会在这里，发生什么了？"
        st.session_state["history"].append(TextMsg({"role": "user", "content": query}))
    input_placeholder.markdown(query)

    response_stream = get_characterglm_response(filter_text_msg(st.session_state["history"]), aNu_character_meta)
    bot_response = output_stream_response(response_stream, message_placeholder)
    if not bot_response:
        message_placeholder.markdown("生成出错")
        st.session_state["history"].pop()
    else:
        st.session_state["history"].append(TextMsg({"role": "user", "content": bot_response}))

def characterglm_LiXiaoYao():
    if st.session_state["history"]:
        query = last_message = st.session_state["history"][-1].get("content")
    else:
        query = "（泪眼，激动的说）逍遥哥哥！你终于醒了！真的是太好了！你知道吗 你昏迷的这几天发生了很多事"
        st.session_state["history"].append(TextMsg({"role": "user", "content": query}))
    input_placeholder.markdown(query)

    response_stream = get_characterglm_response(filter_text_msg(st.session_state["history"]), liXiaoYao_character_meta)
    bot_response = output_stream_response(response_stream, message_placeholder)
    if not bot_response:
        message_placeholder.markdown("生成出错")
        st.session_state["history"].pop()
    else:
        st.session_state["history"].append(TextMsg({"role": "assistant", "content": bot_response}))

def start_chat():
    if not api.API_KEY:
        st.error("未设置API_KEY")
    characterglm_LiXiaoYao()

# 在同一行排列按钮
with st.container():
    n_button = len(button_labels)
    cols = st.columns(n_button)
    button_key_to_col = dict(zip(button_labels.keys(), cols))
    
    with button_key_to_col["start_chat"]:
        if st.button(button_labels["start_chat"], key="start_chat"):
            start_chat()
            st.rerun()
            
    with button_key_to_col["go_on"]:
        if st.button(button_labels["go_on"], key="go_on"):
            if st.session_state["history"] and st.session_state["history"][-1].get("role") == "assistant":
                characterglm_ANu()
            else:
                characterglm_LiXiaoYao()
            st.rerun()
            
    with button_key_to_col["save_chat"]:
        if st.button(button_labels["save_chat"], key="save_chat"):
            save_chat_history()
            st.rerun()
            
    with button_key_to_col["clear_history"]:
        if st.button(button_labels["clear_history"], key="clear_history"):
            init_session()
            st.rerun()

